import torch

# print(torch.__version__)
# 自动求梯度
# torch.autograd.backward(tensors, # 用于求导的张量，如loss
#                         grad_tensors=None, # 如tensors不是标量，需指定grad_tensors，且grad_tensors必须和tensors的维度相同
#                         retain_graph=None, # 保存计算图
#                         create_graph=False) # 创建导数计算图，用于高阶求导
# 参数tensors如果不是标量，需要另外指定参数grad_tensors，参数grad_tensors必须和参数tensors的长度相同。
# 在这一种情况下，backward实际上实现的是代价函数（loss = torch.sum(tensors*grad_tensors);
# 注：torch中向量*向量实际上是点积，因此tensors和grad_tensors的维度必须一致 ）
# 关于叶子节点的梯度计算，而不是参数tensors对于给定图叶子节点的梯度。
# 如果指定参数grad_tensors=torch.ones((size(tensors))),显而易见，代价函数关于叶子节点的梯度，
# 也就等于参数tensors对于给定图叶子节点的梯度。

# 运行两次反向传播
# flag = True
flag = False
if flag:
    w = torch.tensor([1.], requires_grad=True)
    x = torch.tensor([2.], requires_grad=True)

    a = torch.add(w, x)
    b = torch.add(w, 1)
    y = torch.mul(a, b)

    y.backward(create_graph=True)
    print(w.grad)
    y.backward()
    print(w.grad)

# 输出为向量时，计算梯度
# flag = True
flag = False
if flag:
    w = torch.tensor([1.], requires_grad=True)
    x = torch.tensor([2.], requires_grad=True)

    a = torch.add(w, x)
    b = torch.add(w, 1)
    y0 = torch.mul(a, b)  # y0 = (w+x) * (w+1)
    y1 = torch.add(a, b)  # y0 = (w+x) + (w+1)

    # 将y0和y1拼接，输出为 [y0,y1]
    output = torch.cat([y0, y1], dim=0)
    # 输出向量各元素的权重
    grad_tensors = torch.tensor([1., 1.])

    # 相当于 loss = torch.sum(output * grad_tensors)
    output.backward(gradient=grad_tensors)
    print(w.grad)

# 高阶求导
# torch.autograd.grad(outputs, # 求导的张量，因变量，如 loss
#                     inputs, # 求导的张量，自变量，如 x
#                     grad_tensors=None, # 创建导数计算图
#                     retain_graph=None, # 保存导数计算图
#                     create_graph=False) # 多梯度权重
# flag = True
flag = False
if flag:
    x = torch.tensor([3.], requires_grad=True)
    y = torch.pow(x, 2)  # y = x**2

    # 一阶导数
    # 因为要求二阶导数，所以要创建导数计算图
    grad_1 = torch.autograd.grad(y, x, create_graph=True)
    print(grad_1)

    # 二阶导数
    grad_2 = torch.autograd.grad(grad_1, x)
    print(grad_2)

# 原位操作（in-place operation）
# flag = True
flag = False
if flag:

    a = torch.ones((1, ))
    print(id(a), a)

    # = 不是原位操作，操作前后地址变了
    # a = a + torch.ones((1, ))
    # print(id(a), a)

    # += 是原位操作，操作前后地址没有改变
    a += torch.ones((1, ))
    print(id(a), a)

# 叶子结点不能进行原位操作
# 因为反向传播过程中，要用到叶子结点的数据
flag = True
# flag = False
if flag:
    # w和x是叶子结点
    w = torch.tensor([1.], requires_grad=True)
    x = torch.tensor([2.], requires_grad=True)

    a = torch.add(w, x)
    b = torch.add(w, 1)
    y = torch.mul(a, b)

    # 后面带下划线的是原位操作（in-place operation）
    w.add_(1)

    y.backward()